package skewreduce.lib;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Set;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;

import skewreduce.framework.physical.PhysicalOp;


/**
 * new partitioning scheme
 */
public abstract class Partition implements Writable, Comparable<Partition> {
    protected long id;
    protected int lv;         // level -- 0 = root
    protected int ilv;        // inverse level -- 0 = leaf

    protected Partition left;
    protected Partition right;

    public final static class ID implements Comparable<ID>, Writable {
        final long id;
        final int level;

        public ID(String s) {
        	String[] flds = s.split("[-]");
        	level = Integer.parseInt(flds[0]);
        	id = Long.parseLong(flds[1]);
        }
        public ID(long i,int l) { this.id = i; this.level = l; }
        public long getID() { return id; }
        public int getLevel() { return level; }
        
        @Override
        public int hashCode() { return (int)id | (level << 8); }

        @Override
        public boolean equals(Object o) {
            if ( o instanceof ID ) {
                return this == o || (this.id == ((ID)o).id && this.level == ((ID)o).level);
            } else if ( o instanceof PhysicalOp.ID ) {
            	PhysicalOp.ID x = (PhysicalOp.ID)o;
            	return level == x.getLevel() && id == x.getId();
            } else if ( o instanceof Partition ) {
                return id == ((Partition)o).id && level == ((Partition)o).lv;
            }
            return false;
        }
        
        @Override
        public String toString() {
        	return String.format("%d-%d",level,id);
        }
        
		@Override
		public int compareTo(ID o) {
			if ( level != o.level ) return o.level - level;
			return Long.signum(id - o.id);
		}
		
		@Override
		public void readFields(DataInput in) throws IOException {
			throw new UnsupportedOperationException("Use createID() instead");
		}
		@Override
		public void write(DataOutput out) throws IOException {
			out.writeLong(id);
			out.writeInt(level);
		}
		
		public static ID createID(DataInput in) throws IOException {
			long id = in.readLong();
			int lv = in.readInt();
			return new ID(id,lv);
		}
    }
    
    /**
     * semantic is following.
     * compare the two tuples during partitioning.
     * a. in order of fields specified in axes using the values in split. if it is less than the value, go left, otherwise go right.
     * b. if the value is identical to split, then go left according to the leftProb. otherwise go right.
     * 
     * @author yongchul
     *
     */
    public static interface Split extends Writable {
		public int getAxis(int i);
		public float getSplitField(int i);
		public float getLeftProb();
		public void set(String spec);
		public boolean isLeft(ITuple t);
    }

	public static abstract class SplitAdapter implements Split {
		public boolean isLeft(ITuple t) {
			for ( int i = 0; i < t.arity(); ++i ) {
				float v = t.getFloat(getAxis(i));
				float sv = getSplitField(i);
				
				if ( v < sv ) return true;
				if ( v > sv ) return false;
			}
			// it's all same
			return ( getLeftProb() == 0.0 ) ? false : ( Math.random() < getLeftProb() );
		}
	}
    
    protected Partition() {}
    protected Partition(String spec) {
    	set(spec);
    }
    protected Partition(long id,int lv) {
    	this.id = id;
    	this.lv = lv;
    }

    public void set(String spec) {
        String[] fld = spec.split(";");
        if ( fld == null || fld.length <= 2 )
            throw new IllegalArgumentException(spec);

        id = Long.parseLong(fld[0]);
        lv = Integer.parseInt(fld[1]);
        
        setInternal(Arrays.copyOfRange(fld, 2, fld.length));
    }
    protected abstract void setInternal(String[] flds);
    protected abstract Partition newInstance();

    protected abstract void setInternal(Partition org);
    public abstract int getArity();
    public abstract Number getLength(int fld);
    
    // replace the content of this cube
    public synchronized void set(Partition org) {
    	id = org.id;
    	lv = org.lv;
    	left = org.left;
    	right = org.right;
    	setInternal(org);
    }

    public long getID() { return id; }
    public int  getLevel() { return lv; }
    public ID   getPartitionID() { return new ID(id,lv); }

    public boolean isLeaf() { return left == null && right == null; }
    public Partition   getLeft() { return left; }
    public Partition   getRight() { return right; }
    public boolean hasLeft() { return left != null; }
    public boolean hasRight() { return right != null; }
    public abstract double  getVolume();
    
    public void setLeft(Partition p) { left = p; }
    public void setRight(Partition p) { right = p; }

    public boolean isAncestorOf(Partition o) {
    	if ( lv > o.getLevel() ) return false;
    	if ( lv == o.getLevel() )
    		return id == o.getID();
    	// now compare prefix
    	return id == (o.getID() >> (o.getLevel() - lv));
    }
    
    public Partition lookup(ID pid) {
    	// sanity check
    	if ( pid.getLevel() < lv || (pid.getLevel() == lv && pid.getID() != id) ) {
    		return null;
    	}
    	
    	// (lv + 1) bits -- meaningful bits
    	long mask = ~(-1 << (pid.getLevel() - lv));
    	long childBits = pid.getID() & mask; // follow bit pattern until we reach to the child
    	
    	Partition node = this;
    	int diff = pid.getLevel() - lv;
    	while ( diff > 0 ) {
    		long whichChild = ( childBits >>> (diff-1) ) & 0x01L;
    		node = ( whichChild == 0 ) ? node.left : node.right;
    		--diff;
    	}
    	
    	return node;
    }

    /**
     * @return 0 if left, 1 if right. -1 otherwise.
     */
    public int whichChild(Partition o) {
        if ( lv < o.getLevel() && id == (o.getID() >>> 1) ) {
            return (int)(o.getID() & 0x01L);
        }
        return -1;
    }
    
    protected abstract <T extends ITuple> Partition whichPartition(T x);

    public final <T extends ITuple> int whichChild(T x) {
    	return ( whichPartition(x) == left ) ? 0 : 1;
    }
    
    public abstract <T extends ITuple> boolean contains(T p);
	public abstract <T extends ITuple> boolean atMergeSurface(T q,Number eps);

    public abstract <T extends ITuple> boolean atCore(T p,Number eps);
    public <T extends ITuple> boolean atSkin(T p,Number eps) {
        return ! atCore(p,eps);
    }

    public <T extends ITuple> Partition getLeaf(T x) {
        Partition node = this;
        while ( ! node.isLeaf() ) {
        	node = node.whichPartition(x);
        }
        
        assert node.contains(x);
        
        return node;
    }

    // split methods

    public static interface TerminalCondition {
        public boolean continueSplit(int size,int depth);
    }

    public final static class SizeCondition implements TerminalCondition {
        private int sz;
        public SizeCondition(int s) { sz = s; }
        public boolean continueSplit(int size,int depth) {
            return size > sz;
        }
    }

    public final static class DepthCondition implements TerminalCondition {
        private int d;
        public DepthCondition(int depth) { d = depth; }
        public DepthCondition(Partition p,int relDepth) { d = p.getLevel() + relDepth; }
        public boolean continueSplit(int size,int depth) {
            return depth < d;
        }
    }

    public abstract <T extends ITuple> int split(T[] arr,int from,int to,TerminalCondition cond);
    
    public static <T extends ITuple> int build(Partition root,T[] arr,int len,TerminalCondition cond) {
    	return root.split(arr,0,len-1,cond);
    }
    
    // pruning
    
    public void prune(int num) {
        int target = 0;
        while ( num > 1 ) {
            ++target; num >>>= 1;
        }
        prune2(target);
    }

    private void prune2(int target) {
        if ( lv < target ) {
            if ( left != null ) left.prune2(target);
            if ( right != null ) right.prune2(target);
        } else {
            left = null;
            right = null;
        }
    }

    /**
     * @return true when safe to prune
     */
    private boolean prune3(int defLevel,Set<ID> ids) {
        boolean plzPruneL = false;
        boolean plzPruneR = false;
        if ( left != null ) plzPruneL = left.prune3(defLevel,ids);
        if ( right != null ) plzPruneR = right.prune3(defLevel,ids);

        boolean plzPrune = lv > defLevel;

        if ( ids.contains(new ID(id,lv)) ) {
            // make it a leaf node.
            left = null;
            right = null;
            plzPrune = false;
        } else if ( !(plzPruneL && plzPruneR) ) {
            // if any of subtree performed prune, then this should not prune.
            plzPrune = false;
        } else if ( lv == defLevel ) {
            left = null;
            right = null;
            plzPrune = false;
        } else {
        }

        return plzPrune;
    }

    public void prune(int defLevel,Set<ID> ids) {
        prune3(defLevel,ids);
    }

    // writable
    
    public abstract void readInternalFields(DataInput in) throws IOException;

    public void readFields(DataInput in) throws IOException {
        lv = in.readByte();
        id = in.readLong();
        
        readInternalFields(in);

        int childflag = in.readByte();
        left  = ( (childflag & 0x01) == 0 ) ? null : read(in);
        right = ( (childflag & 0x02) == 0 ) ? null : read(in);
    }

    public abstract void writeInternalFields(DataOutput out) throws IOException;
    
    public void write(DataOutput out) throws IOException {
        out.writeByte(lv);
        out.writeLong(id);
        writeInternalFields(out);

        int flag = ( left == null ) ? 0 : 0x01;
        flag |= (right == null ) ? 0 : 0x02;
        out.writeByte(flag);
        if ( left != null ) left.write(out);
        if ( right != null ) right.write(out);
    }

    public Partition read(DataInput in) throws IOException {
        Partition r = newInstance();
        r.readFields(in);
        return r;
    }

    @SuppressWarnings("unchecked")
	public static <T extends Partition> T read(Configuration conf,DataInput in) throws IOException {
    	Class<T> cls = (Class<T>)conf.getClass("skewreduce.partition.class", Partition.class);
    	return read(cls,in);
    }
    
    public static <T extends Partition> T read(Class<T> cls,DataInput in) throws IOException {
    	try {
			T r = cls.newInstance();
			r.readFields(in);
			return r;
		} catch (InstantiationException e) {
			throw new IOException(e);
		} catch (IllegalAccessException e) {
			throw new IOException(e);
		}
    }

    // print interface
    public void print(java.io.PrintStream out) {
        java.io.PrintWriter wr = new java.io.PrintWriter(out);
    	print(wr);
        wr.flush();
    }
    public void print(java.io.PrintWriter out) {
        if ( left != null ) left.print(out);
        if ( right != null ) right.print(out);
        out.println(this);
    }

    protected void toSpecInternal(StringBuilder buf) {
    }
    
    public String toSpec() {
        StringBuilder buf = new StringBuilder();
        buf.append(id);
        buf.append(';');
        buf.append(lv);
        buf.append(';');
        
        toSpecInternal(buf);

        return buf.toString();
    }

    protected String toStringInternal() {
    	return "";
    }
    
    @Override
    public String toString() {
        return String.format("%d/%d;%s;",lv,id,toStringInternal());
    }
    
    public String toLongString() {
    	java.io.StringWriter wr = new java.io.StringWriter(8192);
    	java.io.PrintWriter out = new java.io.PrintWriter(wr);
    	print(out);
    	out.flush();
    	out.close();
    	return wr.getBuffer().toString();
    }

    @Override
    public final int hashCode() { return (int)(lv * id); }

    @Override
    public boolean equals(Object o) {
        if ( this == o ) return true;
        if ( o instanceof Partition ) {
            return id == ((Partition)o).id && lv == ((Partition)o).lv;
        }
        return false;
    }

	public int uniformSplit(TerminalCondition cond) {
		throw new UnsupportedOperationException();
	}
	
	@Override
	public int compareTo(Partition o) {
		if ( lv != o.lv ) return o.lv - lv;
		return Long.signum(id - o.id);
	}
	/**
	 * deep copy subtree.
	 * @return
	 * @throws IOException
	 */
	public Partition deepCopy() throws IOException {
		ByteArrayOutputStream baos = null;
		DataOutputStream out = null;
		ByteArrayInputStream bais = null;
		DataInputStream in = null;

		try {
			baos = new ByteArrayOutputStream(65536);
			out = new DataOutputStream(baos);
			this.write(out);
			out.flush();
			byte[] content = baos.toByteArray();
			
			bais = new ByteArrayInputStream(content);
			in = new DataInputStream(bais);
			return Partition.read(this.getClass(), in);
		} finally {
			if ( out != null ) try { out.close(); } catch ( IOException ignore ) {}
			out = null;
			if ( in != null ) try { in.close(); } catch ( IOException ignore ) {}
			in = null;
		}
	}
	
	public Partition shallowCopy() {
		Partition clone = newInstance();
		clone.set(this);
		return clone;
	}
	
	public Partition[] split(int axis,Number v) {
		return split(axis,v,0.0f);
	}
	public abstract Partition[] split(int axis,Number v,float lprob);
	public abstract Partition[] split(int[] axis,float[] v,float lprob); // split multiple axes
	public abstract Partition getMergingArea(Partition l,Partition r,Number eps);
	
	public void dump(String fn) throws IOException {
		FileOutputStream fout = null;
		DataOutputStream dout = null;
		try {
			fout = new FileOutputStream(fn);
			dout = new DataOutputStream(fout);
			
			write(dout);
			
		} catch ( IOException ex ) {
			throw ex;
		} finally {
			if ( dout != null ) try { dout.close(); } catch ( IOException ignore ) {}
			dout = null;
			if ( fout != null ) try { fout.close(); } catch ( IOException ignore ) {}
			fout = null;
		}
	}
}
